import os
import json
import zipfile
import hashlib
import shutil
from datetime import datetime
from typing import Dict, List, Optional, Tuple
import difflib
import argparse

class ScratchGit:
    def __init__(self, project_path: str):
        self.project_path = os.path.abspath(project_path)
        self.scratchgit_dir = os.path.join(self.project_path, '.scratchgit')
        self.objects_dir = os.path.join(self.scratchgit_dir, 'objects')
        self.refs_dir = os.path.join(self.scratchgit_dir, 'refs')
        self.head_file = os.path.join(self.scratchgit_dir, 'HEAD')
        self.index_file = os.path.join(self.scratchgit_dir, 'INDEX')
        
        self._init_dirs()
    
    def _init_dirs(self):
        dirs_to_create = [
            self.scratchgit_dir,
            self.objects_dir,
            self.refs_dir,
            os.path.join(self.refs_dir, 'heads'),
            os.path.join(self.refs_dir, 'remotes')
        ]
        
        for dir_path in dirs_to_create:
            if not os.path.exists(dir_path):
                os.makedirs(dir_path)
    
    def _hash_object(self, data: bytes) -> str:
        return hashlib.sha1(data).hexdigest()
    
    def _store_object(self, data: bytes) -> str:
        obj_hash = self._hash_object(data)
        obj_path = os.path.join(self.objects_dir, obj_hash)
        
        if not os.path.exists(obj_path):
            with open(obj_path, 'wb') as f:
                f.write(data)
        
        return obj_hash
    
    def _read_object(self, obj_hash: str) -> bytes:
        obj_path = os.path.join(self.objects_dir, obj_hash)
        if not os.path.exists(obj_path):
            raise FileNotFoundError(f"Object {obj_hash} not found")
        
        with open(obj_path, 'rb') as f:
            return f.read()
    
    def _get_current_branch(self) -> str:
        if not os.path.exists(self.head_file):
            return 'master'
        
        with open(self.head_file, 'r') as f:
            content = f.read().strip()
            if content.startswith('ref: refs/heads/'):
                return content.split('/')[-1]
        return 'master'
    
    def _get_current_commit(self) -> Optional[str]:
        branch = self._get_current_branch()
        branch_file = os.path.join(self.refs_dir, 'heads', branch)
        
        if os.path.exists(branch_file):
            with open(branch_file, 'r') as f:
                return f.read().strip()
        return None
    
    def _set_current_commit(self, commit_hash: str, branch: Optional[str] = None):
        if branch is None:
            branch = self._get_current_branch()
        
        branch_file = os.path.join(self.refs_dir, 'heads', branch)
        with open(branch_file, 'w') as f:
            f.write(commit_hash)
    
    def _create_commit_object(self, tree_hash: str, parent_hash: Optional[str], 
                            message: str, author: str) -> str:
        commit_data = {
            'tree': tree_hash,
            'parent': parent_hash,
            'author': author,
            'timestamp': datetime.now().isoformat(),
            'message': message
        }
        
        commit_json = json.dumps(commit_data, indent=2, ensure_ascii=False)
        return self._store_object(commit_json.encode('utf-8'))
    
    def _create_tree_object(self, project_data: dict, assets_hashes: Dict[str, str]) -> str:
        tree_data = {
            'project': project_data,
            'assets': assets_hashes
        }
        
        tree_json = json.dumps(tree_data, indent=2, ensure_ascii=False)
        return self._store_object(tree_json.encode('utf-8'))
    
    def _extract_sb3(self, sb3_path: str) -> Tuple[dict, Dict[str, bytes]]:

        project_data = {}
        assets = {}
        
        with zipfile.ZipFile(sb3_path, 'r') as zip_ref:
            for file_info in zip_ref.filelist:
                if file_info.filename == 'project.json':
                    with zip_ref.open(file_info) as f:
                        project_data = json.loads(f.read())
                else:
                    with zip_ref.open(file_info) as f:
                        assets[file_info.filename] = f.read()
        
        return project_data, assets
    
    def _create_sb3(self, project_data: dict, assets: Dict[str, bytes], output_path: str):
        with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zip_ref:
            zip_ref.writestr('project.json', json.dumps(project_data, indent=2, ensure_ascii=False))
            
            for filename, data in assets.items():
                zip_ref.writestr(filename, data)
    
    def _get_commit_data(self, commit_hash: str) -> dict:
        return json.loads(self._read_object(commit_hash))
    
    def _get_tree_data(self, tree_hash: str) -> dict:
        return json.loads(self._read_object(tree_hash))
    
    def init(self):
        self._init_dirs()
        print(f"Initialized empty ScratchGit repository in {self.scratchgit_dir}")
    
    def add(self, sb3_path: str):
        if not os.path.exists(sb3_path):
            raise FileNotFoundError(f"File {sb3_path} not found")
        
        project_data, assets = self._extract_sb3(sb3_path)
        
        assets_hashes = {}
        for filename, data in assets.items():
            asset_hash = self._store_object(data)
            assets_hashes[filename] = asset_hash
        
        project_json = json.dumps(project_data, indent=2, ensure_ascii=False)
        project_hash = self._store_object(project_json.encode('utf-8'))
        
        tree_hash = self._create_tree_object(project_data, assets_hashes)
        
        index_data = {
            'tree': tree_hash,
            'project_hash': project_hash,
            'assets_hashes': assets_hashes
        }
        
        with open(self.index_file, 'w', encoding='utf-8') as f:
            json.dump(index_data, f, indent=2, ensure_ascii=False)
        
        print(f"Added {sb3_path} to staging area")
    
    def commit(self, message: str, author: str = "Anonymous"):
        if not os.path.exists(self.index_file):
            print("Nothing to commit, working directory clean")
            return
        
        with open(self.index_file, 'r', encoding='utf-8') as f:
            index_data = json.load(f)
        
        tree_hash = index_data['tree']
        parent_hash = self._get_current_commit()
        
        commit_hash = self._create_commit_object(tree_hash, parent_hash, message, author)
        
        self._set_current_commit(commit_hash)
        
        os.remove(self.index_file)
        
        print(f"[{commit_hash[:7]}] {message}")
        return commit_hash
    
    def log(self, limit: int = 10, branch: Optional[str] = None):
        if branch is None:
            current_commit = self._get_current_commit()
        else:
            branch_file = os.path.join(self.refs_dir, 'heads', branch)
            if os.path.exists(branch_file):
                with open(branch_file, 'r') as f:
                    current_commit = f.read().strip()
            else:
                print(f"Branch {branch} not found")
                return
        
        if not current_commit:
            print("No commits yet")
            return
        
        commits = []
        commit_hash = current_commit
        
        for i in range(limit):
            try:
                commit_data = self._get_commit_data(commit_hash)
                commits.append({
                    'hash': commit_hash,
                    'message': commit_data['message'],
                    'author': commit_data['author'],
                    'timestamp': commit_data['timestamp'],
                    'parent': commit_data['parent']
                })
                
                if commit_data['parent']:
                    commit_hash = commit_data['parent']
                else:
                    break
            except Exception as e:
                break
        
        for commit in commits:
            print(f"commit {commit['hash']}")
            print(f"Author: {commit['author']}")
            print(f"Date: {commit['timestamp']}")
            print(f"    {commit['message']}")
            print()
    
    def diff(self):
        current_commit = self._get_current_commit()
        if not current_commit:
            print("No commits yet")
            return
        
        commit_data = self._get_commit_data(current_commit)
        tree_data = self._get_tree_data(commit_data['tree'])
        last_project = tree_data['project']
        
        sb3_files = [f for f in os.listdir(self.project_path) if f.endswith('.sb3')]
        if not sb3_files:
            print("No .sb3 file found in current directory")
            return
        
        current_sb3 = os.path.join(self.project_path, sb3_files[0])
        current_project, _ = self._extract_sb3(current_sb3)
        
        last_json = json.dumps(last_project, indent=2, sort_keys=True)
        current_json = json.dumps(current_project, indent=2, sort_keys=True)
        
        if last_json != current_json:
            print("Project code differences:")
            diff = difflib.unified_diff(
                last_json.splitlines(keepends=True),
                current_json.splitlines(keepends=True),
                fromfile='HEAD',
                tofile='working directory'
            )
            print(''.join(diff))
        else:
            print("No differences in project code")
    
    def checkout(self, commit_or_branch: str, output_path: Optional[str] = None):
        branch_file = os.path.join(self.refs_dir, 'heads', commit_or_branch)
        if os.path.exists(branch_file):
            with open(branch_file, 'r') as f:
                commit_hash = f.read().strip()
        else:
            commit_hash = commit_or_branch
        
        try:
            commit_data = self._get_commit_data(commit_hash)
            
            tree_data = self._get_tree_data(commit_data['tree'])
            
            project_data = tree_data['project']
            
            assets = {}
            for filename, asset_hash in tree_data['assets'].items():
                assets[filename] = self._read_object(asset_hash)
            
            # 创建sb3文件
            if output_path is None:
                output_path = os.path.join(self.project_path, f"checkout_{commit_hash[:7]}.sb3")
            
            self._create_sb3(project_data, assets, output_path)
            print(f"Checked out commit {commit_hash[:7]} to {output_path}")
            
        except Exception as e:
            print(f"Error checking out commit: {e}")
    
    def status(self):
        current_commit = self._get_current_commit()
        branch = self._get_current_branch()
        
        print(f"On branch {branch}")
        
        if current_commit:
            print(f"Current commit: {current_commit[:7]}")
        else:
            print("No commits yet")
        
        if os.path.exists(self.index_file):
            print("Changes to be committed:")
            print("  (use 'scratchgit commit' to commit)")
            print("        new file: project.json and assets")
        else:
            print("Nothing to commit, working directory clean")
    
    def branch(self, branch_name: Optional[str] = None, start_point: Optional[str] = None):
        if branch_name is None:
            heads_dir = os.path.join(self.refs_dir, 'heads')
            if os.path.exists(heads_dir):
                branches = os.listdir(heads_dir)
                current_branch = self._get_current_branch()
                for branch in branches:
                    if branch == current_branch:
                        print(f"* {branch}")
                    else:
                        print(f"  {branch}")
            return
        
        if start_point is None:
            start_point = self._get_current_commit()
        
        if start_point is None:
            print("No commits yet, cannot create branch")
            return
        
        try:
            self._get_commit_data(start_point)
        except:
            branch_file = os.path.join(self.refs_dir, 'heads', start_point)
            if os.path.exists(branch_file):
                with open(branch_file, 'r') as f:
                    start_point = f.read().strip()
            else:
                print(f"Invalid start point: {start_point}")
                return
        
        branch_file = os.path.join(self.refs_dir, 'heads', branch_name)
        with open(branch_file, 'w') as f:
            f.write(start_point)
        
        print(f"Created branch '{branch_name}' pointing to {start_point[:7]}")
    
    def checkout_branch(self, branch_name: str):
        branch_file = os.path.join(self.refs_dir, 'heads', branch_name)
        if not os.path.exists(branch_file):
            print(f"Branch '{branch_name}' not found")
            return
        
        # 更新HEAD文件指向新分支
        with open(self.head_file, 'w') as f:
            f.write(f'ref: refs/heads/{branch_name}')
        
        print(f"Switched to branch '{branch_name}'")

    def merge(self, branch_name: str, message: Optional[str] = None):
        self._enhanced_merge(branch_name, message)

    def stash(self, message: str = "WIP"):
        stash_dir = os.path.join(self.scratchgit_dir, 'stash')
        if not os.path.exists(stash_dir):
            os.makedirs(stash_dir)
        
        current_commit = self._get_current_commit()
        branch = self._get_current_branch()
        
        stash_data = {
            'commit': current_commit,
            'branch': branch,
            'message': message,
            'timestamp': datetime.now().isoformat(),
            'staged': None
        }
        
        if os.path.exists(self.index_file):
            with open(self.index_file, 'r') as f:
                stash_data['staged'] = json.load(f)
            os.remove(self.index_file)
        
        stash_file = os.path.join(stash_dir, f"stash_{int(datetime.now().timestamp())}")
        with open(stash_file, 'w') as f:
            json.dump(stash_data, f, indent=2)
        
        print(f"Saved working directory to stash: {message}")

    def stash_list(self):
        stash_dir = os.path.join(self.scratchgit_dir, 'stash')
        if not os.path.exists(stash_dir):
            print("No stashes")
            return
        
        stashes = []
        for filename in os.listdir(stash_dir):
            if filename.startswith('stash_'):
                stash_file = os.path.join(stash_dir, filename)
                with open(stash_file, 'r') as f:
                    stash_data = json.load(f)
                stashes.append((filename, stash_data))
        
        if not stashes:
            print("No stashes")
            return
        
        stashes.sort(key=lambda x: x[0])
        for filename, data in stashes:
            timestamp = data.get('timestamp', 'Unknown')
            message = data.get('message', 'No message')
            print(f"{filename}: {message} [{timestamp}]")

    def stash_pop(self):
        stash_dir = os.path.join(self.scratchgit_dir, 'stash')
        if not os.path.exists(stash_dir):
            print("No stashes")
            return

        stashes = []
        for filename in os.listdir(stash_dir):
            if filename.startswith('stash_'):
                stashes.append(filename)
        
        if not stashes:
            print("No stashes")
            return
        
        stashes.sort()
        latest_stash = stashes[-1]
        stash_file = os.path.join(stash_dir, latest_stash)
        
        with open(stash_file, 'r') as f:
            stash_data = json.load(f)
        
        if stash_data.get('staged'):
            with open(self.index_file, 'w') as f:
                json.dump(stash_data['staged'], f, indent=2)
            print("Restored staged changes")

        os.remove(stash_file)
        print(f"Popped stash: {latest_stash}")


    def merge(self, branch_name: str, message: Optional[str] = None):
        current_branch = self._get_current_branch()
        if current_branch == branch_name:
            print("Cannot merge branch with itself")
            return
        
        current_commit = self._get_current_commit()
        branch_file = os.path.join(self.refs_dir, 'heads', branch_name)
        
        if not os.path.exists(branch_file):
            print(f"Branch '{branch_name}' not found")
            return
        
        with open(branch_file, 'r') as f:
            target_commit = f.read().strip()
        
        if not current_commit:
            print("No commits in current branch")
            return
        
        ancestor = self._find_merge_base(current_commit, target_commit)
        if not ancestor:
            print("No common ancestor found")
            return
        
        if ancestor == current_commit:
            self._set_current_commit(target_commit)
            print(f"Fast-forward merge: {current_commit[:7]} -> {target_commit[:7]}")
            return
        
        if ancestor == target_commit:
            print("Already up to date")
            return
        
        print(f"Merging {branch_name} into {current_branch}")
        
        ancestor_tree = self._get_tree_data(self._get_commit_data(ancestor)['tree'])
        current_tree = self._get_tree_data(self._get_commit_data(current_commit)['tree'])
        target_tree = self._get_tree_data(self._get_commit_data(target_commit)['tree'])

        merged_project, conflicts = self._merge_project_data(
            ancestor_tree['project'],
            current_tree['project'],
            target_tree['project']
        )
        
        merged_assets = self._merge_assets(
            ancestor_tree['assets'],
            current_tree['assets'],
            target_tree['assets']
        )

        if conflicts:
            print("Merge conflicts detected:")
            for conflict in conflicts:
                print(f"  - {conflict}")
            print("Please resolve conflicts manually and commit")
            return
        
        # 存储合并后的数据
        assets_hashes = {}
        for filename, data in merged_assets.items():
            asset_hash = self._store_object(data)
            assets_hashes[filename] = asset_hash
        
        tree_hash = self._create_tree_object(merged_project, assets_hashes)
        
        if message is None:
            message = f"Merge branch '{branch_name}'"
        
        merge_commit = self._create_commit_object(
            tree_hash, 
            current_commit, 
            message, 
            "ScratchGit"
        )
        
        self._set_current_commit(merge_commit)
        print(f"Merge commit created: {merge_commit[:7]}")
    
    def _find_merge_base(self, commit1: str, commit2: str) -> Optional[str]:
        ancestors1 = self._get_ancestors(commit1)
        ancestors2 = self._get_ancestors(commit2)

        for ancestor in ancestors1:
            if ancestor in ancestors2:
                return ancestor
        return None
    
    def _get_ancestors(self, commit_hash: str) -> List[str]:
        ancestors = []
        current = commit_hash
        
        while current:
            ancestors.append(current)
            try:
                commit_data = self._get_commit_data(current)
                current = commit_data['parent']
            except:
                break
        
        return ancestors
    
    def _merge_project_data(self, ancestor: dict, current: dict, target: dict) -> Tuple[dict, List[str]]:
        conflicts = []
        merged = target.copy()
        
        ancestor_json = json.dumps(ancestor, sort_keys=True)
        current_json = json.dumps(current, sort_keys=True)
        target_json = json.dumps(target, sort_keys=True)
        
        if ancestor_json != current_json and ancestor_json != target_json:
            if current_json != target_json:
                conflicts.append("project.json")
        
        return merged, conflicts
    
    def _merge_assets(self, ancestor: Dict[str, str], current: Dict[str, str], 
                     target: Dict[str, str]) -> Dict[str, bytes]:
        merged_assets = {}

        all_files = set(ancestor.keys()) | set(current.keys()) | set(target.keys())
        
        for filename in all_files:
            ancestor_hash = ancestor.get(filename)
            current_hash = current.get(filename)
            target_hash = target.get(filename)
            
            if target_hash:
                merged_assets[filename] = self._read_object(target_hash)
            elif current_hash:
                merged_assets[filename] = self._read_object(current_hash)
            elif ancestor_hash:
                merged_assets[filename] = self._read_object(ancestor_hash)
        
        return merged_assets

def main():
    parser = argparse.ArgumentParser(description='ScratchGit - Git-like version control for Scratch projects')
    parser.add_argument('command', nargs='?', choices=[
        'init', 'add', 'commit', 'log', 'diff', 'checkout', 'status',
        'branch', 'checkout-branch', 'merge', 'stash'
    ], help='Command to execute')
    parser.add_argument('args', nargs='*', help='Command arguments')
    parser.add_argument('--message', '-m', help='Commit message')
    parser.add_argument('--author', help='Author name')
    parser.add_argument('--output', '-o', help='Output file path for checkout')
    
    args = parser.parse_args()
    
    project_path = os.getcwd()
    
    # 创建ScratchGit实例
    sg = ScratchGit(project_path)
    
    try:
        if args.command is None:
            parser.print_help()
            return
        
        if args.command == 'init':
            sg.init()
        
        elif args.command == 'add':
            if not args.args:
                print("Usage: scratchgit add <sb3_file>")
                return
            sg.add(args.args[0])
        
        elif args.command == 'commit':
            if not args.message:
                print("Usage: scratchgit commit -m 'commit message'")
                return
            author = args.author or "Anonymous"
            sg.commit(args.message, author)
        
        elif args.command == 'log':
            limit = int(args.args[0]) if args.args and args.args[0].isdigit() else 10
            branch = args.args[1] if len(args.args) > 1 else None
            sg.log(limit, branch)
        
        elif args.command == 'diff':
            sg.diff()
        
        elif args.command == 'checkout':
            if not args.args:
                print("Usage: scratchgit checkout <commit_or_branch>")
                return
            sg.checkout(args.args[0], args.output)
        
        elif args.command == 'status':
            sg.status()
        
        elif args.command == 'branch':
            if args.args:
                branch_name = args.args[0]
                start_point = args.args[1] if len(args.args) > 1 else None
                sg.branch(branch_name, start_point)
            else:
                sg.branch()
        
        elif args.command == 'checkout-branch':
            if not args.args:
                print("Usage: scratchgit checkout-branch <branch_name>")
                return
            sg.checkout_branch(args.args[0])
        
        elif args.command == 'merge':
            if not args.args:
                print("Usage: scratchgit merge <branch_name>")
                return
            sg.merge(args.args[0], args.message)
        
        elif args.command == 'stash':
            if args.args and args.args[0] == 'list':
                sg.stash_list()
            elif args.args and args.args[0] == 'pop':
                sg.stash_pop()
            else:
                message = args.message or "WIP"
                sg.stash(message)
    
    except Exception as e:
        print(f"Error: {e}")
    
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    main()
